package com.aloudmeta.grpc.support.metric;

import io.grpc.ForwardingServerCall;
import io.grpc.Metadata;
import io.grpc.ServerCall;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.micrometer.core.instrument.Counter;

/**
 * huaixin 2022/2/9 2:08 PM
 */
class MetricCollectingServerCall<Q, A> extends ForwardingServerCall.SimpleForwardingServerCall<Q, A> {

    private final Counter responseCounter;
    private Code responseCode = Code.UNKNOWN;

    /**
     * Creates a new delegating ServerCall that will wrap the given server call to collect metrics.
     *
     * @param delegate The original call to wrap.
     * @param responseCounter The counter for incoming responses.
     */
    public MetricCollectingServerCall(
            final ServerCall<Q, A> delegate,
            final Counter responseCounter) {

        super(delegate);
        this.responseCounter = responseCounter;
    }

    public Code getResponseCode() {
        return this.responseCode;
    }

    @Override
    public void close(final Status status, final Metadata responseHeaders) {
        this.responseCode = status.getCode();
        super.close(status, responseHeaders);
    }

    @Override
    public void sendMessage(final A responseMessage) {
        this.responseCounter.increment();
        super.sendMessage(responseMessage);
    }

}
